package com.samp.solr.update;

import org.apache.http.client.HttpClient;
import org.apache.solr.client.solrj.impl.ConcurrentUpdateSolrClient;
import org.apache.solr.client.solrj.impl.HttpClientUtil;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.params.ModifiableSolrParams;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

/**
 * 测试各种提交索引的性能。
 */
public class UpdateTest {

    private final static String SOLR_URL = "http://localhost:8080/solr/samp";

    private static List<DocDTO> getDocList(){
        String filePath = "solr/data_update.txt";
        URL url = UpdateTest.class.getClassLoader().getResource(filePath);
        File file = new File(url.getFile());
        BufferedReader reader = null;
        List<DocDTO> solrData = new ArrayList<>();
        try {
            reader = new BufferedReader(new FileReader(file));
            DocDTO docDTO = new DocDTO();
            int index = 0 ;
            String tempString ;
            while ((tempString = reader.readLine()) != null){
                if( index == 0 ){
                    docDTO.setTitle(tempString.trim());
                }else if( index == 1 ){
                    docDTO.setContent(tempString.trim());
                }
                index++;
            }
//            System.out.println(docDTO.getTitle());
//            System.out.println(docDTO.getContent());
            solrData.add(docDTO);
        }catch(Exception e) {
            e.printStackTrace();
        }finally {
            if (reader != null) {
                try {
                    reader.close();
                }catch(IOException e1){
                    e1.printStackTrace();
                }
            }
        }
        return solrData;
    }

    private static void commitWithSolrj() throws Exception{
        HttpSolrClient solr = new HttpSolrClient.Builder(SOLR_URL).build();
        Long start = System.currentTimeMillis();
        List<DocDTO> docList = getDocList();
        for( int i = 0 ; i < 1000000; i++ ){
            for(DocDTO doc: docList){
                SolrInputDocument document = new SolrInputDocument();
                document.addField("id", UUID.randomUUID().toString());
                document.addField("title", doc.getTitle());
                document.addField("content", doc.getContent());
                solr.add(document);
//                solr.add(document,10000);
//                solr.commit(true,true,true);
            }
        }
        System.out.println("cost:"+(System.currentTimeMillis()-start));
    }

    private static void commitWithConcurrentUpdateSolrClient()throws Exception{
        ModifiableSolrParams params = new ModifiableSolrParams();
        params.set(HttpClientUtil.PROP_MAX_CONNECTIONS, 128);
        params.set(HttpClientUtil.PROP_MAX_CONNECTIONS_PER_HOST, 32);
        params.set(HttpClientUtil.PROP_FOLLOW_REDIRECTS, false);
        params.set(HttpClientUtil.PROP_ALLOW_COMPRESSION, true);
        params.set(HttpClientUtil.PROP_USE_RETRY, true);
        HttpClient httpClient = HttpClientUtil.createClient(params);
        ConcurrentUpdateSolrClient updateSolrClient = new ConcurrentUpdateSolrClient.Builder(SOLR_URL)
                            .withHttpClient(httpClient).withQueueSize(100).withThreadCount(10).build();
        updateSolrClient.setConnectionTimeout(1000);
//        updateSolrClient.setSoTimeout(1000);
        updateSolrClient.setPollQueueTime(500);
        Long start = System.currentTimeMillis();
        List<DocDTO> docList = getDocList();
        for( int i = 0 ; i < 1000000; i++ ){
            for(DocDTO doc: docList){
                SolrInputDocument document = new SolrInputDocument();
                document.addField("id", UUID.randomUUID().toString());
                document.addField("title", doc.getTitle());
                document.addField("content", doc.getContent());
                updateSolrClient.add(document);
            }
        }
        System.out.println("cost:"+(System.currentTimeMillis()-start));
    }

    private static void deleteAllDoc() throws Exception{
        HttpSolrClient solrClient = new HttpSolrClient.Builder(SOLR_URL).build();
        solrClient.deleteByQuery("id:*");
        solrClient.commit();
    }

    public static void main(String[] args) {
        try {
//            deleteAllDoc();
//            commitWithSolrj();
            commitWithConcurrentUpdateSolrClient();
        }catch (Exception e){
            e.printStackTrace();
        }
        System.out.println("done!");
    }
}

class DocDTO{
    private String title;
    private String content;

    public String getTitle() {
        return title;
    }

    public void setTitle(String title) {
        this.title = title;
    }

    public String getContent() {
        return content;
    }

    public void setContent(String content) {
        this.content = content;
    }
}
